#include <iostream>
#include <connection/udp.hpp>


namespace udp_proxy{
ConnUdp::ConnUdp(std::string bind_host, uint16_t bind_port,
               std::string remote_host, uint16_t remote_port){
    _isopen = false;

    _buf_ptr = 0;
    _buf_len = 0;

    _local_host = bind_host;
    _local_port = bind_port;
    _remote_host = remote_host;
    _remote_port = remote_port;

    // std::cout << bind_host << ":" << bind_port << std::endl;
    _setup();

#if defined(__linux__)
    _socket = -1;
    // Start mutex
    int result = pthread_mutex_init(&_lock, NULL);
    if (result != 0) {
        printf("ERROR: mutex init failed\n");
        throw 1;
    }
#endif //
}

ConnUdp::~ConnUdp(){
    close();
}

int ConnUdp::connect(const received_cb cb){
    _received_cb = cb;
    _thread_recv = std::thread(std::bind(&ConnUdp::_thread_recv_callback, this));

    return 0;
}


void ConnUdp::_thread_recv_callback(){
    uint8_t buf[512];
    int len;
    while(!_should_exit){
        len = receive_data(buf);
        _received_cb(buf, len);

        // printf("received len=%d\n", len);
        // std::this_thread::sleep_for(std::chrono::milliseconds(1000));
    }
}

size_t ConnUdp::receive_data(uint8_t *buf){
    int bytes_recv;

#if defined(__linux__)
    // Lock
    // pthread_mutex_lock(&_lock);
    socklen_t addr_len = sizeof(struct sockaddr_in);
    bytes_recv         = recvfrom(_socket, (void *)buf, RECEIVE_BUF_MAX_SIZE, 0, (struct sockaddr *)&_dest_addr, &addr_len);

    // Unlock
    // pthread_mutex_unlock(&_lock);

#elif defined(_WIN32)
    int addr_len = sizeof(SOCKADDR_IN);
    bytes_recv   = recvfrom(_socket, (char *)buf, RECEIVE_BUF_MAX_SIZE, 0, (SOCKADDR *)&_dest_addr, &addr_len);
#endif //
    // printf("INFO: received size=%d\n", bytes_recv);
    return bytes_recv;
}

size_t ConnUdp::send_data(uint8_t *buf, size_t len){
    int bytes_send;

    std::lock_guard<std::mutex> lock(_tx_mutex);
    // printf("sending data, len=%d\n", len);

#if defined(__linux__)
    bytes_send = sendto(_socket, buf, len, 0, (struct sockaddr *)&_dest_addr, sizeof(struct sockaddr_in));
#elif defined(_WIN32)
    bytes_send   = sendto(_socket, (char *)buf, len, 0, (SOCKADDR *)&_dest_addr, sizeof(SOCKADDR));
#endif //

    return bytes_send;
}


void ConnUdp::close(){
    int ret;

    _should_exit = true;
    if (this->_isopen) {
#if defined(__linux__)
        ret     = ::close(_socket);
        _socket = -1;
#elif defined(_WIN32)
        ret = closesocket(_socket);
        WSACleanup();
#endif //
        if (ret) {
            fprintf(stderr, "WARNING: Error on udp port close (%i)\n", ret);
        } else {
            printf("INFO: close udp port success\n");
        }

        _isopen = false;
    }
    _thread_recv.join();
    printf("ConnUdp::close()");
}


bool ConnUdp::_setup() {
#if defined(__linux__)
    _socket = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (_socket < 0) {
        printf("Error: sock init failed\n");
        throw EXIT_FAILURE;
    }
    /* local addr */
    memset(&_loc_addr, 0, sizeof(_loc_addr));
    _loc_addr.sin_family      = AF_INET;
    _loc_addr.sin_addr.s_addr = INADDR_ANY;
    _loc_addr.sin_port        = htons(_local_port);
    /* destination addr */
    memset(&_dest_addr, 0, sizeof(_dest_addr));
    _dest_addr.sin_family      = AF_INET;
    _dest_addr.sin_addr.s_addr = inet_addr(_remote_host.c_str());
    _dest_addr.sin_port        = htons(_remote_port);

    // 这里我们将网络接收设置为非阻塞，即如果没有数据，判断超时后将跳出接收继续执行而不是一直阻塞在接收
    // 设置非阻塞，方式一：
    // int flags  = fcntl(_socket, F_GETFL, 0);
    // int result = fcntl(_socket, F_SETFL, flags | O_NONBLOCK);
    // if (result == -1) {
    //     perror("ERROR: set UDP none block failed\n");
    // }

    // 设置非阻塞，方式二：
    /* You must do this for:
     * 1. recvfrom() will not blocked if you set timeout, otherwise recvfrom() will
     * blocked while not data receiveing.
     * 2. without it, you cann't exit program by Ctrl+C
     */
    struct timeval timeout;
    timeout.tv_sec  = 0;
    timeout.tv_usec = 1000;
    if (setsockopt(_socket, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) == -1) {
        perror("ERROR: setsocketopt failed.\n");
        return false;
    }

    /* bind the socket */
    if (bind(_socket, (struct sockaddr *)&_loc_addr, sizeof(struct sockaddr))) {
        perror("ERROR: sock bind failed\n");
        ::close(_socket);
        _socket = -1;
        return false;
    }
#elif defined(_WIN32)
    WSADATA WSAData;

    if (WSAStartup(MAKEWORD(2, 2), &WSAData) != 0) {
        printf("ERROR: at WSAStartup()\n");
        exit(1);
    }
    _socket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (_socket == INVALID_SOCKET) {
        printf("ERROR: Failed socket()\n");
        return 0;
    }
    // 本地地址
    _loc_addr.sin_family           = AF_INET;
    _loc_addr.sin_addr.S_un.S_addr = INADDR_ANY;
    _loc_addr.sin_port             = htons(_local_port);
    // 设置非阻塞接收
    // unsigned long mode   = 1;
    // int           result = ::ioctlsocket(_socket, FIONBIO, &mode);
    // if (result != 0) {
    //     printf("ERROR: set UDP none block failed\n");
    // }
    // UDP服务器绑定
    if (bind(_socket, (SOCKADDR *)&_loc_addr, sizeof(SOCKADDR)) == SOCKET_ERROR) {
        printf("ERROR: sock bind failed with error %u\n", WSAGetLastError());
        return false;
    }
    // 远端地址
    _dest_addr.sin_family           = AF_INET;
    _dest_addr.sin_addr.S_un.S_addr = inet_addr(_remote_host.c_str());
    _dest_addr.sin_port             = htons(_remote_port);

#endif //
    return true;
}
}